import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import OneHotEncoder


class CustomFeatureProcessor(BaseEstimator, TransformerMixin):
    def __init__(self, skew_threshold=0.5):
        """
        自定义特征处理器

        参数:
            skew_threshold (float): 判断数值特征是否需要对数变换的偏态阈值
        """
        self.skew_threshold = skew_threshold
        self.nominal_features = ['BusinessTravel', 'Department', 'EducationField',
                                 'Gender', 'JobRole', 'MaritalStatus', 'OverTime']
        self.encoder = None
        self.skewed_features = None
        self.numeric_cols = None

    def fit(self, X, y=None):
        """拟合数据：学习编码器和偏态特征"""
        # 1. 初始化独热编码器
        self.encoder = OneHotEncoder(handle_unknown='ignore', sparse_output=False)
        self.encoder.fit(X[self.nominal_features])

        # 2. 检测数值特征的偏态
        self.numeric_cols = X.select_dtypes(include=['float64', 'int64']).columns
        skewness = X[self.numeric_cols].skew()
        self.skewed_features = skewness[abs(skewness) > self.skew_threshold].index.tolist()

        return self

    def transform(self, X):
        """应用转换：编码类别特征和处理偏态"""
        # 创建副本避免修改原始数据
        X_transformed = X.copy()

        # 1. 独热编码类别特征
        encoded = self.encoder.transform(X_transformed[self.nominal_features])
        encoded_df = pd.DataFrame(
            encoded,
            columns=self.encoder.get_feature_names_out(self.nominal_features),
            index=X_transformed.index
        )
        X_transformed = X_transformed.drop(self.nominal_features, axis=1)
        X_transformed = pd.concat([X_transformed, encoded_df], axis=1)

        # 2. 处理数值特征的偏态
        if self.skewed_features:
            for col in self.skewed_features:
                X_transformed[col] = np.log1p(X_transformed[col])

        return X_transformed

    def get_feature_names_out(self, input_features=None):
        """获取输出特征名称（用于Pipeline兼容性）"""
        # 获取数值特征名称（处理过偏态的和未处理的）
        numeric_features = [f for f in self.numeric_cols if f not in self.nominal_features]

        # 获取编码后的类别特征名称
        encoded_features = self.encoder.get_feature_names_out(self.nominal_features)

        return np.array(numeric_features + list(encoded_features))